import os
import shutil
import logging
import threading
from queue import Queue
from datetime import datetime
import multiprocessing

def setup_logging():
    """设置日志记录"""
    log_dir = "file_operation_logs"
    os.makedirs(log_dir, exist_ok=True)
    
    log_filename = datetime.now().strftime(f"{log_dir}/file_operation_%Y%m%d_%H%M%S.log")
    logging.basicConfig(
        filename=log_filename,
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        encoding='utf-8'
    )
    return log_filename

def get_folder_path(prompt):
    """获取并验证文件夹路径"""
    while True:
        path = input(prompt).strip()
        # 处理引号包围的路径
        if path.startswith(('"', "'")) and path.endswith(('"', "'")):
            path = path[1:-1]
        
        if os.path.isdir(path):
            return os.path.abspath(path)
        else:
            print(f"错误：'{path}' 不是一个有效的文件夹路径。请重新输入。")

def get_source_folders():
    """获取用户输入的多个源文件夹"""
    source_folders = []
    print("\n请输入要从中操作文件的源文件夹路径。")
    print("输入完一个路径后按回车，输入空行结束。")
    
    while True:
        folder = input("源文件夹路径: ").strip()
        if not folder:
            break
            
        # 处理引号包围的路径
        if folder.startswith(('"', "'")) and folder.endswith(('"', "'")):
            folder = folder[1:-1]
            
        if os.path.isdir(folder):
            abs_path = os.path.abspath(folder)
            if abs_path not in source_folders:
                source_folders.append(abs_path)
                print(f"已添加: {abs_path}")
            else:
                print("这个文件夹已经添加过了。")
        else:
            print(f"警告：'{folder}' 不是一个有效的文件夹，已跳过。")
    
    return source_folders

def get_operation_mode():
    """获取用户选择的操作模式（移动或复制）"""
    print("\n===== 操作模式选择 =====")
    while True:
        choice = input("请选择操作模式 (1-移动文件, 2-复制文件): ").strip()
        if choice in ('1', '2'):
            return 'move' if choice == '1' else 'copy'
        else:
            print("无效选择，请输入 1 或 2")

def get_file_filters():
    """获取用户设置的文件过滤条件"""
    filters = {
        'include_extensions': [],
        'exclude_extensions': [],
        'include_keywords': [],
        'exclude_keywords': []
    }
    
    print("\n===== 文件过滤设置 =====")
    use_filter = input("是否要设置文件过滤条件？(y/n): ").strip().lower()
    if use_filter not in ('y', 'yes'):
        return None
    
    # 扩展名过滤
    include_ext = input("只包含这些扩展名（用逗号分隔，如: txt,jpg）: ").strip().lower()
    if include_ext:
        filters['include_extensions'] = [f".{ext.strip()}" if not ext.startswith('.') else ext.strip() 
                                        for ext in include_ext.split(',')]
    
    exclude_ext = input("排除这些扩展名（用逗号分隔）: ").strip().lower()
    if exclude_ext:
        filters['exclude_extensions'] = [f".{ext.strip()}" if not ext.startswith('.') else ext.strip() 
                                        for ext in exclude_ext.split(',')]
    
    # 关键词过滤
    include_key = input("文件名包含这些关键词（用逗号分隔）: ").strip().lower()
    if include_key:
        filters['include_keywords'] = [key.strip().lower() for key in include_key.split(',')]
    
    exclude_key = input("文件名排除这些关键词（用逗号分隔）: ").strip().lower()
    if exclude_key:
        filters['exclude_keywords'] = [key.strip().lower() for key in exclude_key.split(',')]
    
    return filters

def should_include_file(filename, filters):
    """根据过滤条件判断是否应包含该文件"""
    if not filters:
        return True
        
    lower_filename = filename.lower()
    
    # 检查扩展名包含
    if filters['include_extensions']:
        has_matching_ext = any(lower_filename.endswith(ext) for ext in filters['include_extensions'])
        if not has_matching_ext:
            return False
    
    # 检查扩展名排除
    if filters['exclude_extensions']:
        has_excluded_ext = any(lower_filename.endswith(ext) for ext in filters['exclude_extensions'])
        if has_excluded_ext:
            return False
    
    # 检查关键词包含
    if filters['include_keywords']:
        has_matching_key = any(key in lower_filename for key in filters['include_keywords'])
        if not has_matching_key:
            return False
    
    # 检查关键词排除
    if filters['exclude_keywords']:
        has_excluded_key = any(key in lower_filename for key in filters['exclude_keywords'])
        if has_excluded_key:
            return False
    
    return True

def collect_files_to_process(source_folders, dest_folder, recursive, filters):
    """收集所有要处理的文件信息"""
    files_to_process = []
    
    for source in source_folders:
        source_abs = os.path.abspath(source)
        dest_abs = os.path.abspath(dest_folder)
        
        # 避免源文件夹和目标文件夹相同
        if source_abs == dest_abs:
            print(f"跳过源文件夹 '{source}'，它与目标文件夹相同。")
            logging.warning(f"跳过源文件夹 '{source}'，它与目标文件夹相同。")
            continue
        
        # 递归遍历文件夹
        walk_generator = os.walk(source) if recursive else [(source, [], os.listdir(source))]
        
        for root, _, files in walk_generator:
            for file in files:
                file_path = os.path.join(root, file)
                
                # 检查是否为文件
                if not os.path.isfile(file_path):
                    continue
                
                # 应用过滤条件
                if not should_include_file(file, filters):
                    continue
                
                # 计算目标路径
                dest_path = os.path.join(dest_folder, file)
                
                # 处理文件名冲突
                counter = 1
                temp_dest = dest_path
                while os.path.exists(temp_dest):
                    name, ext = os.path.splitext(file)
                    temp_dest = os.path.join(dest_folder, f"{name}_{counter}{ext}")
                    counter += 1
                
                # 获取文件大小用于优化复制策略
                file_size = os.path.getsize(file_path)
                files_to_process.append((file_path, temp_dest, file_size))
    
    return files_to_process

def preview_files(files_to_process, operation_mode):
    """预览将要处理的文件"""
    operation_name = "移动" if operation_mode == 'move' else "复制"
    print(f"\n===== {operation_name}预览 =====")
    
    if not files_to_process:
        print("没有找到符合条件的可处理文件")
        return False
        
    # 显示前10个文件
    for src, dest, _ in files_to_process[:10]:
        print(f"{src} -> {dest}")
    
    # 如果文件数量超过10个，显示省略信息
    if len(files_to_process) > 10:
        print(f"... 还有 {len(files_to_process)-10} 个文件将被{operation_name}")
    
    total_size = sum(size for _, _, size in files_to_process)
    total_size_mb = total_size / (1024 * 1024)
    print(f"\n总计 {len(files_to_process)} 个文件 ({total_size_mb:.2f} MB) 将被{operation_name}")
    return True

def copy_file_optimized(src, dest, file_size):
    """优化的文件复制函数，根据文件大小动态调整缓冲区"""
    # 根据文件大小选择最佳缓冲区大小
    if file_size < 1024 * 1024:  # 小于1MB的小文件
        block_size = 64 * 1024  # 64KB
    elif file_size < 100 * 1024 * 1024:  # 1MB到100MB的中等文件
        block_size = 512 * 1024  # 512KB
    else:  # 大于100MB的大文件
        block_size = 4 * 1024 * 1024  # 4MB
    
    try:
        with open(src, 'rb') as fsrc, open(dest, 'wb') as fdest:
            # 使用os.sendfile（如果系统支持）实现零拷贝
            if hasattr(os, 'sendfile') and file_size > 0:
                fsrc_fd = fsrc.fileno()
                fdest_fd = fdest.fileno()
                offset = 0
                while offset < file_size:
                    sent = os.sendfile(fdest_fd, fsrc_fd, offset, min(block_size, file_size - offset))
                    if sent == 0:
                        break
                    offset += sent
            else:
                # 回退到常规缓冲复制
                while True:
                    buf = fsrc.read(block_size)
                    if not buf:
                        break
                    fdest.write(buf)
        
        # 复制文件元数据
        shutil.copystat(src, dest)
        return True, None
    except Exception as e:
        # 复制失败时清理目标文件
        if os.path.exists(dest):
            try:
                os.remove(dest)
            except:
                pass
        return False, str(e)

def worker(queue, operation_mode, results, lock):
    """多线程工作函数"""
    while not queue.empty():
        item = queue.get()
        if item is None:
            break
            
        src, dest, file_size = item
        success = False
        error = None
        
        try:
            if operation_mode == 'move':
                shutil.move(src, dest)
                success = True
            else:  # copy
                success, error = copy_file_optimized(src, dest, file_size)
        except Exception as e:
            success = False
            error = str(e)
        
        # 线程安全地更新结果
        with lock:
            if success:
                results['success'] += 1
                logging.info(f"成功{'移动' if operation_mode == 'move' else '复制'}: {src} -> {dest}")
            else:
                results['failed'] += 1
                error_msg = f"{'移动' if operation_mode == 'move' else '复制'}失败 '{src}': {error or '未知错误'}"
                print(f"\n{error_msg}")
                logging.error(error_msg)
        
        queue.task_done()

def process_files(files_to_process, operation_mode, use_tqdm):
    """使用多线程执行文件处理操作（移动或复制）"""
    operation_name = "移动" if operation_mode == 'move' else "复制"
    total = len(files_to_process)
    
    if total == 0:
        print("没有文件需要处理")
        return 0, 0
    
    # 确定最佳线程数（通常为CPU核心数的1-2倍，复制操作可适当增加）
    cpu_count = multiprocessing.cpu_count()
    thread_count = min(cpu_count * 2, total) if operation_mode == 'copy' else min(cpu_count, total)
    
    print(f"使用 {thread_count} 个线程进行{operation_name}操作...")
    
    # 初始化线程安全的队列和结果计数器
    queue = Queue()
    for item in files_to_process:
        queue.put(item)
    
    results = {'success': 0, 'failed': 0}
    lock = threading.Lock()
    threads = []
    
    # 创建并启动线程
    for _ in range(thread_count):
        t = threading.Thread(
            target=worker,
            args=(queue, operation_mode, results, lock)
        )
        t.daemon = True
        t.start()
        threads.append(t)
    
    # 显示进度
    if use_tqdm:
        from tqdm import tqdm
        with tqdm(total=total, desc=f"{operation_name}文件中", unit="个文件") as pbar:
            while results['success'] + results['failed'] < total:
                current = results['success'] + results['failed']
                pbar.update(current - pbar.n)
            pbar.update(total - pbar.n)
    else:
        print(f"{operation_name}文件中...")
        while results['success'] + results['failed'] < total:
            progress = results['success'] + results['failed']
            if progress % 10 == 0 or progress == total:
                print(f"已处理 {progress}/{total} 个文件")
    
    # 等待所有线程完成
    for t in threads:
        t.join()
    
    print(f"\n操作完成。成功{operation_name} {results['success']} 个文件，失败 {results['failed']} 个文件。")
    logging.info(f"操作完成。成功{operation_name} {results['success']} 个文件，失败 {results['failed']} 个文件。")
    
    return results['success'], results['failed']

def main():
    print("===== 高效文件移动/复制工具 =====")
    print("本工具采用多线程和动态优化技术，可快速处理大量文件。")
    
    # 检查tqdm是否可用
    try:
        import tqdm
        use_tqdm = True
        print("检测到tqdm，将使用进度条显示")
    except ImportError:
        use_tqdm = False
        print("未检测到tqdm，将使用基础进度显示")
        print("提示: 安装tqdm可以获得更好的进度显示体验 (pip install tqdm)")
    
    # 设置日志
    log_file = setup_logging()
    print(f"操作日志将保存到: {log_file}")
    
    # 获取操作模式（移动或复制）
    operation_mode = get_operation_mode()
    operation_name = "移动" if operation_mode == 'move' else "复制"
    print(f"已选择{operation_name}模式")
    
    # 获取源文件夹列表
    source_folders = get_source_folders()
    
    if not source_folders:
        print("没有添加任何源文件夹，程序退出。")
        return
    
    print("\n已选择的源文件夹:")
    for i, folder in enumerate(source_folders, 1):
        print(f"{i}. {folder}")
    
    # 获取递归选项
    recursive = input("\n是否递归处理子文件夹中的文件？(y/n): ").strip().lower() in ('y', 'yes')
    
    # 获取文件过滤条件
    filters = get_file_filters()
    
    # 获取目标文件夹
    dest_folder = get_folder_path(f"\n请输入目标文件夹路径（文件将被{operation_name}到这里）: ")
    
    # 确保目标文件夹存在
    os.makedirs(dest_folder, exist_ok=True)
    
    # 收集要处理的文件（包含文件大小信息）
    print(f"\n正在扫描文件（将{operation_name}的文件）...")
    files_to_process = collect_files_to_process(source_folders, dest_folder, recursive, filters)
    
    if not files_to_process:
        print(f"没有找到可{operation_name}的文件，程序退出。")
        return
    
    # 预览模式
    preview = input(f"\n是否要预览将要{operation_name}的文件？(y/n): ").strip().lower() in ('y', 'yes')
    if preview:
        if not preview_files(files_to_process, operation_mode):
            print("程序退出。")
            return
    
    # 确认操作
    total_size = sum(size for _, _, size in files_to_process)
    total_size_mb = total_size / (1024 * 1024)
    print(f"\n即将{operation_name} {len(files_to_process)} 个文件 ({total_size_mb:.2f} MB) 到: {dest_folder}")
    confirm = input("是否继续? (y/n): ").strip().lower()
    
    if confirm in ('y', 'yes'):
        process_files(files_to_process, operation_mode, use_tqdm)
        print(f"操作日志已保存到: {log_file}")
    else:
        print("操作已取消。")

if __name__ == "__main__":
    main()
    